import tensorflow as tf
import numpy as np
from data_util import  TextData
from train import Encoder, Decoder
from args import Args


def evaluate(sentence):
    attention_plot = np.zeros((max_length, max_length))

    sentence = textData.sen2enco(sentence)

    inputs = tf.keras.preprocessing.sequence.pad_sequences([sentence],
                                                           maxlen=max_length,
                                                           padding='post')
    inputs = tf.convert_to_tensor(inputs)

    result = ''

    hidden = [tf.zeros((1, units))]
    enc_out, enc_hidden = encoder(inputs, hidden)

    dec_hidden = enc_hidden
    dec_input = tf.expand_dims([1], 0)

    for t in range(max_length):
        predictions, dec_hidden, attention_weights = decoder(dec_input,
                                                             dec_hidden,
                                                             enc_out)

        # 存储注意力权重以便后面制图
        attention_weights = tf.reshape(attention_weights, (-1, ))
        attention_plot[t] = attention_weights.numpy()

        predicted_id = tf.argmax(predictions[0]).numpy()

        if textData.sr_id2word[predicted_id] != "<eos>":
            result += textData.sr_id2word[predicted_id]
        else:
            return result, attention_plot

        # 预测的 ID 被输送回模型
        dec_input = tf.expand_dims([predicted_id], 0)

    return result, attention_plot


if __name__ == '__main__':
    args = Args()    
    textData = TextData(args)
    
    BATCH_SIZE = 512
    embedding_dim = 64
    units = 512
    vocab_size = textData.vocab_size
    max_length = 18
    
    encoder = Encoder(vocab_size, embedding_dim, units, BATCH_SIZE)
    decoder = Decoder(vocab_size, embedding_dim, units, BATCH_SIZE)
    
    optimizer = tf.keras.optimizers.Adam()
    checkpoint = tf.train.Checkpoint(optimizer=optimizer,
                                     encoder=encoder,
                                     decoder=decoder)
    checkpoint.restore(tf.train.latest_checkpoint("./save_model/seq2seq"))   
    
    while True:
        a = input("input:")
        result, _ = evaluate(a)
        print(result)
    